package com.zbkj.client.config;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.zbkj.client.controller.MessageWebSocketHandler;
import com.zbkj.client.response.WebsocketResponse;
import com.zbkj.common.result.CommonResult;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Configuration;
import org.springframework.http.HttpStatus;
import org.springframework.web.socket.config.annotation.EnableWebSocket;
import org.springframework.web.socket.config.annotation.WebSocketConfigurer;
import org.springframework.web.socket.config.annotation.WebSocketHandlerRegistry;
import org.springframework.web.socket.server.support.HttpSessionHandshakeInterceptor;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.server.HandshakeInterceptor;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.http.server.ServletServerHttpRequest;

import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.Map;

@Configuration
@EnableWebSocket
public class WebSocketConfig implements WebSocketConfigurer {
    @Autowired
    private MessageWebSocketHandler messageWebSocketHandler;

    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
        registry.addHandler(messageWebSocketHandler, "/ws")
                .setAllowedOrigins("*")
                .addInterceptors(new WebSocketTokenInterceptor());
    }
    
    /**
     * WebSocket Token 拦截器
     * 用于处理 WebSocket 连接时的 token 验证
     */
    public static class WebSocketTokenInterceptor implements HandshakeInterceptor {
        private final ObjectMapper objectMapper = new ObjectMapper();
        
        @Override
        public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response,
                                       WebSocketHandler wsHandler, Map<String, Object> attributes) throws Exception {
            if (request instanceof ServletServerHttpRequest) {
                ServletServerHttpRequest servletRequest = (ServletServerHttpRequest) request;
                
                // 从请求头中获取token
                String token = servletRequest.getServletRequest().getHeader("token");
                
                // 如果请求头中没有，尝试从URL参数中获取（兼容性）
                if (token == null || token.trim().isEmpty()) {
                    token = servletRequest.getServletRequest().getParameter("token");
                }
                
                if (token != null && !token.trim().isEmpty()) {
                    // 将 token 保存到 WebSocket 会话属性中
                    attributes.put("token", token);
                    
                    // 这里可以添加 token 验证逻辑
                    // 例如：验证 token 是否有效，获取用户信息等
                    System.out.println("WebSocket 连接，Token: " + token);
                    
                    return true;
                } else {
                    System.out.println("WebSocket 连接缺少 Token");
                    // 设置401状态码
                    response.setStatusCode(HttpStatus.UNAUTHORIZED);

                    // 返回指定JSON
                    String jsonResponse = "{\"event\":\"Unauthorized\",\"data\":\"未登录\"}";
                    try {
                        response.getBody().write(jsonResponse.getBytes(StandardCharsets.UTF_8));
                    } catch (Exception e) {
                        e.printStackTrace();
                    }
                    
                    return false; // 拒绝连接
                }
            }
            return true;
        }
        
        @Override
        public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response,
                                  WebSocketHandler wsHandler, Exception exception) {
            // 握手完成后的处理
        }
    }
} 